#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import argparse
import numpy as np
import time


def args_parser():
    parser = argparse.ArgumentParser(prog="Unweighted DiGraph Generator")
    # source/target --vertex --edge --number --output
    parser.add_argument('type', choices=['single', 'multiple'], help="choose whether to generate single graph or multiple graphs")
    parser.add_argument('-v', '--vertex', type=int, dest='vertex', help="the number of vertex (Required)", required=True)
    parser.add_argument('-e', '--edge', type=int, dest='edge', help='the number of edge   (Required)', required=True)
    parser.add_argument('-n', '--number', type=int, dest='number', help='the number of graph you want to generate',
                        default=1)
    parser.add_argument('-o', '--output', dest='output', help='the output file name')
    args = parser.parse_args()

    # Ensure: vertex < edge < vertex * (vertex - 1) / 2
    if args.edge > (args.vertex * (args.vertex - 1) / 2):
        print('[ERROR] the number of edge is more than (vertex * (vertex - 1) / 2)')
        exit(1)
    if args.edge < args.vertex:
        print('[ERROR] the number of edge is too small,'
              ' it cannot be less than the number of vertex because we must ensure the digraph is connected')
        exit(1)

    return args.type, args.vertex, args.edge, args.number, args.output

def generate_single_graph(vertex_num, edge_num, output):
    output = 'V{}E{}.txt'.format(vertex_num, edge_num) if not output else output
    f = open(output, 'w')
    print('# Directed graph (each unordered pair of nodes is saved once): {}'.format(output), file=f)
    print('# Generated by unweighted-digraph-generator.py on {}'.format(
        time.strftime("%Y-%m-%d %H:%M:%S %A", time.localtime())), file=f)
    print('# Vertexes: {} Edges: {}'.format(vertex_num, edge_num), file=f)
    print('# FromVertexId	ToVertexId', file=f)

    vertex_array = np.arange(vertex_num)
    np.random.shuffle(vertex_array)
    edge_set = set()
    for _ in range(edge_num):
        edge = vertex_array[np.random.randint(vertex_num, size=2)]
        stamp = str(np.sort(edge))
        while stamp in edge_set or edge[0] == edge[1]: # exclude self-loop
            edge = vertex_array[np.random.randint(vertex_num, size=2)]
            stamp = str(np.sort(edge))
        edge_set.add(stamp)
        print('{}   {}'.format(edge[0], edge[1]), file=f)

    print('the generate graph have been save in: {}.'.format(output))
    f.close()

def generate_multiple_graphs(vertex_num, edge_num, number, output):
    output = 'Q{}-{}.txt'.format(vertex_num, number) if not output else output
    f = open(output, 'w')
    vertex_array = np.arange(vertex_num)
    for n in range(number):
        np.random.shuffle(vertex_array)

        print('t # {}'.format(n), file=f)
        for v in range(vertex_num):
            print('v {}'.format(v), file=f)

        edge_set = set()
        for _ in range(edge_num):
            edge = vertex_array[np.random.randint(vertex_num, size=2)]
            stamp = str(np.sort(edge))
            while stamp in edge_set or edge[0] == edge[1]: # exclude self-loop
                edge = vertex_array[np.random.randint(vertex_num, size=2)]
                stamp = str(np.sort(edge))
            edge_set.add(stamp)
            print('e {} {}'.format(edge[0], edge[1]), file=f)
    print('t # -1', file=f)
    
    print('the generate graph have been save in: {}.'.format(output))
    f.close()

if __name__ == '__main__':
    _type, vertex_num, edge_num, number, output = args_parser()

    if _type == 'single':
        if number > 1:
            print("[WARN} '-n --number' doesn't work for 'single', it's been ignored.")
        generate_single_graph(vertex_num, edge_num, output)
    else:
        generate_multiple_graphs(vertex_num, edge_num, number, output)
